package com.gitee.wsl.struct.tree.visitor

import com.gitee.wsl.struct.tree.NodeElement
import com.gitee.wsl.struct.tree.Tree
import com.gitee.wsl.struct.tree.TreeElement

/**
 * Interface for AST visitors.
 */
interface TreeVisitor<T,E: TreeElement<T, E>> {
    /**
     * Callback executed when entering a node during the visit. The visited node
     * is provided as a parameter of the callback.
     */
    fun startTree(tree: E)

    /**
     * Callback executed when exiting a node during the visit. The visited node
     * is provided as a parameter of the callback.
     */
    fun endTree(tree: E)

    /**
     * A default does nothing visitor.
     */
    open class DefaultTreeVisitor<T,E: TreeElement<T, E>> : TreeVisitor<T,E> {
        override fun startTree(tree: E) {}
        override fun endTree(tree: E) {}
    }

    /**
     * A does nothing visitor that distinguished between visiting an inner node
     * and visiting a leaf node.
     */
   open class InnerNodesAndLeavesVisitor<T,E: TreeElement<T, E>> : TreeVisitor<T,E> {
        /**
         * {@inheritDoc}
         * This method is final, users have to use the following methods to customize
         * their visitor.
         * @see .startInnerNode
         * @see .visitLeaf
         */
        override fun startTree(tree: E) {
            if (tree.isLeaf) visitLeaf(tree) else startInnerNode(tree)
        }

        /**
         * {@inheritDoc}
         * This method is final, users have to use the following methods to customize
         * their visitor.
         * @see .endInnerNode
         * @see .visitLeaf
         */
        override fun endTree(tree: E) {
            if (!tree.isLeaf) endInnerNode(tree)
        }

        /**
         * Callback executed when entering an inner node during the visit. The visited node
         * is provided as a parameter of the callback.
         */
        fun startInnerNode(tree: E) {}

        /**
         * Callback executed when visiting a leaf node during the visit. The visited node
         * is provided as a parameter of the callback.
         */
        fun visitLeaf(tree: E) {}

        /**
         * Callback executed when entering an inner node during the visit. The visited node
         * is provided as a parameter of the callback.
         */
        fun endInnerNode(tree: E) {}
    }

    companion object {
        /**
         * Start visiting the given tree using the given visitor.
         * The tree is visited in post-order.
         */
        fun<T,E: TreeElement<T, E>> visitTree(root: E, visitor: TreeVisitor<T,E>) {
            val stack:ArrayDeque<Pair<E, Iterator<E>>> = ArrayDeque()
            if(root is NodeElement<*,*> && !root.isLeaf)
               stack.add(root to root.children!!.iterator() as Iterator<E>)
            visitor.startTree(root)
            while (stack.isNotEmpty()) {
                val (first, second) = stack.removeAt(0)
                if (!second.hasNext()) {
                    visitor.endTree(first)
                    stack.removeFirst()
                } else {
                    val child = second.next()
                    if(child is NodeElement<*,*> && !root.isLeaf)
                        stack.add(child to child.children!!.iterator() as Iterator<E>)
                    visitor.startTree(child)
                }
            }
        }
    }
}


fun<T,E: TreeElement<T, E>> Tree<T, E>.visitor(visitor: TreeVisitor<T,E>){
    children.forEach {
        TreeVisitor.visitTree(it,visitor)
    }
}